#include "box_LU.hpp"


template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
BoxLU<idx_t, data_t, setup_t, calc_t, dof>::BoxLU() {}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
BoxLU<idx_t, data_t, setup_t, calc_t, dof>::~BoxLU()
{
    if (A_calc && own_A) { delete    A_calc   ; A_calc    = nullptr; }
    delete locA_csr; locA_csr = nullptr;
    delete aux; aux = nullptr;
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void BoxLU<idx_t, data_t, setup_t, calc_t, dof>::truncate() {}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void BoxLU<idx_t, data_t, setup_t, calc_t, dof>::SetOperator(const Operator<idx_t, setup_t, dof> & op)
{
    assert(this->setup_called == false);
    const par_CSRMatrix<idx_t, setup_t, setup_t, dof> & par_A = (const par_CSRMatrix<idx_t, setup_t, setup_t, dof> &) op;
    int my_pid; MPI_Comm_rank(par_A.comm, & my_pid);
    if (my_pid == 0) printf("Setup Box_LU #dof %d\n", dof);
    this->setup_time = - wall_time();

    if constexpr (sizeof(setup_t) == sizeof(calc_t) && sizeof(setup_t) == sizeof(data_t)) {
        own_A = false;
        A_calc = & par_A;
    }
    else {
        own_A = true;
        par_CSRMatrix<idx_t, data_t, calc_t, dof> * A_copy = new par_CSRMatrix<idx_t, data_t, calc_t, dof>
            (par_A.comm, par_A.rows_partition, par_A.cols_partition);
        // 拷贝数据，但不拷贝通信包
        A_copy->commpkg           = par_A.commpkg;
        A_copy->commpkg_transpose = par_A.commpkg_transpose;
        A_copy->col_map_offd      = par_A.col_map_offd;
        auto copy_trunc = [](const CSRMatrix<idx_t, setup_t, setup_t, dof> & src, CSRMatrix<idx_t, data_t, calc_t, dof> & dst) {
            dst.nrows = src.nrows;
            dst.ncols = src.ncols;
            dst.nnz   = src.nnz;
            dst.alloc_mem(false);

            #pragma omp parallel
            {
                #pragma omp for schedule(static) nowait
                for (idx_t i = 0; i <= dst.nrows; i++) {
                    dst.row_ptr[i] = src.row_ptr[i];
                }
                #pragma omp for schedule(static) nowait
                for (idx_t j = 0; j < dst.nnz; j++) {
                    dst.col_idx[j] = src.col_idx[j];
                }
                #pragma omp for schedule(static) nowait
                for (idx_t j = 0; j < dst.nnz * (long long) dof*dof; j++) {
                    dst.vals   [j] = src.vals   [j];
                }
            }
        };
        copy_trunc(par_A.diag, A_copy->diag);
        copy_trunc(par_A.offd, A_copy->offd);

        A_calc = A_copy;
    }

    // 然后对对角块做LU分解
    factorize(par_A.diag);
    // 中间结果
    aux = new seq_Vector<idx_t, calc_t, dof>(A_calc->diag.nrows);

    this->setup_called = true;
    this->setup_time += wall_time();
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void BoxLU<idx_t, data_t, setup_t, calc_t, dof>::Mult_ipl(const par_Vector<idx_t, calc_t, dof> & input, par_Vector<idx_t, calc_t, dof> & output) const
{
    if (this->zero_guess) {
        // 直接执行前代回代
        locA_csr->apply(input.local_vector->data, output.local_vector->data, 0);
    }
    else {
        par_Vector<idx_t, calc_t, dof> glb_resi(input);
        // resi <- A*x
        this->A_calc->Mult(output, glb_resi, false);
        vec_add(input, (calc_t) -1.0, glb_resi, glb_resi);
        // 对残差resi执行前代回代
        locA_csr->apply(glb_resi.local_vector->data, aux->data, 0);
        seq_vec_add<idx_t, calc_t, calc_t, dof>(*(output.local_vector), (calc_t) 1.0, *aux, *(output.local_vector));
    }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void BoxLU<idx_t, data_t, setup_t, calc_t, dof>::factorize(const CSRMatrix<idx_t, setup_t, setup_t, dof> & A)
{
    assert(A.nrows == A.ncols);
    if constexpr (dof == 1) {// 输入的A矩阵已经是一个符合直接法求解器输入的格式
        calc_t* dof1_val = new calc_t [A.nnz];// 注意这些传给CSR_sparseMat类的数组不需要在此处释放！！！
        #pragma omp parallel for schedule(static)
        for (idx_t j = 0; j < A.nnz; j++)
            dof1_val[j] = A.vals[j];

        locA_csr = new CSR_sparseMat<idx_t, calc_t, calc_t>(0, A.nrows, A.row_ptr, A.col_idx, dof1_val);
    } else { static_assert(dof > 1); // 需要将BSR格式转换成CSR矩阵
        // 注意这些传给CSR_sparseMat类的数组不需要在此处释放！！！
        idx_t * dof1_rpt = nullptr;
        idx_t * dof1_cid = nullptr;
        calc_t* dof1_val = nullptr;

        constexpr int e_size = dof*dof;

        dof1_rpt = new idx_t [A.nrows * dof + 1]; 
        dof1_rpt[0] = 0;
        
        #pragma omp parallel for schedule(static)
        for (idx_t bi = 0; bi < A.nrows; bi ++) {
            idx_t cnts[dof];
            for (int f = 0; f < dof; f++) cnts[f] = 0;

            for (auto p = A.row_ptr[bi]; p < A.row_ptr[bi + 1]; p ++) {
                const setup_t * bv = A.vals + p * e_size;
                for (int r = 0; r < dof; r++)
                for (int c = 0; c < dof; c++) {
#ifdef COL_MAJOR
                    const setup_t cv = bv[c * dof + r];
#else // ROW_MAJOR
                    const setup_t cv = bv[r * dof + c];
#endif
                    if (cv != 0.0) cnts[r] ++;
                }
            }
            for (int f = 0; f < dof; f++)
                dof1_rpt[bi * dof + f + 1] = cnts[f];
        }// bi loop

        for (idx_t ci = 0; ci < A.nrows * dof; ci ++)
            dof1_rpt[ci + 1] += dof1_rpt[ci];
        
        const idx_t csr_nnz = dof1_rpt[A.nrows * dof];

        dof1_cid = new idx_t  [csr_nnz];
        dof1_val = new calc_t [csr_nnz];

        #pragma omp parallel for schedule(static)
        for (idx_t bi = 0; bi < A.nrows; bi ++) {
            idx_t cnts[dof];
            for (int f = 0; f < dof; f++) cnts[f] = 0;
            const idx_t * offsets = dof1_rpt + bi * dof;

            for (auto p = A.row_ptr[bi]; p < A.row_ptr[bi + 1]; p ++) {
                const idx_t bj = A.col_idx[p];
                const setup_t * bv = A.vals + p * e_size;
                for (int r = 0; r < dof; r++)
                for (int c = 0; c < dof; c++) {
                    const idx_t cj = bj * dof + c;
#ifdef COL_MAJOR
                    const setup_t cv = bv[c * dof + r];
#else // ROW_MAJOR
                    const setup_t cv = bv[r * dof + c];
#endif
                    if (cv != 0.0) {
                        dof1_cid[offsets[r] + cnts[r]] = cj;
                        dof1_val[offsets[r] + cnts[r]] = cv;
                        cnts[r] ++;
                    }
                }
            }
            for (int f = 0; f < dof; f++)
                assert(cnts[f] + offsets[f] == offsets[f + 1]);
        }// bi loop

        locA_csr = new CSR_sparseMat<idx_t, calc_t, calc_t>(0, A.nrows * dof, dof1_rpt, dof1_cid, dof1_val);
    }

    bool succ;
    succ = locA_csr->initA(0, CSR_sparseMat<idx_t, calc_t, calc_t>::PATTERN::NONSYMMETRIC); assert(succ);
    succ = locA_csr->initB(1, 0); assert(succ);
    succ = locA_csr->initX(1, 0); assert(succ);
    succ = locA_csr->initSolver(omp_get_max_threads(), 1, 0); assert(succ);
    // 符号分析时间不计入??
    this->setup_time += wall_time();
    locA_csr->analyse(1, 0, 1, omp_get_max_threads(), 1);
    this->setup_time -= wall_time();
    // 数值分解
    locA_csr->factorize(1e-4);
}


template class BoxLU<int, double, double, double, 1>;
template class BoxLU<int, double, double, double, 2>;
template class BoxLU<int, double, double, double, 3>;
template class BoxLU<int, double, double, double, 4>;

template class BoxLU<int, float , double, double, 1>;
template class BoxLU<int, float , double, double, 2>;
template class BoxLU<int, float , double, double, 3>;
template class BoxLU<int, float , double, double, 4>;

template class BoxLU<int, float , double, float , 1>;
template class BoxLU<int, float , double, float , 2>;
template class BoxLU<int, float , double, float , 3>;
template class BoxLU<int, float , double, float , 4>;

template class BoxLU<int, float , float , float , 1>;
template class BoxLU<int, float , float , float , 2>;
template class BoxLU<int, float , float , float , 3>;
template class BoxLU<int, float , float , float , 4>;
#ifdef USE_FP16
template class BoxLU<int, __fp16, double, float , 1>;
template class BoxLU<int, __fp16, double, float , 2>;
template class BoxLU<int, __fp16, double, float , 3>;
template class BoxLU<int, __fp16, double, float , 4>;

template class BoxLU<int, __fp16, float , float , 1>;
template class BoxLU<int, __fp16, float , float , 2>;
template class BoxLU<int, __fp16, float , float , 3>;
template class BoxLU<int, __fp16, float , float , 4>;
#endif